package com.ziwei.material.interceptor;

import brave.Span;
import brave.propagation.ThreadLocalSpan;
import com.mysql.cj.MysqlConnection;
import com.mysql.cj.Query;
import com.mysql.cj.interceptors.QueryInterceptor;
import com.mysql.cj.jdbc.JdbcConnection;
import com.mysql.cj.log.Log;
import com.mysql.cj.protocol.Resultset;
import com.mysql.cj.protocol.ServerSession;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URI;
import java.sql.SQLException;
import java.util.Properties;
import java.util.function.Supplier;

/**
 * description:使用zipkin实现mysql执行跟踪
 *
 * @author shengshenglalala
 * @date 2021/2/1 17:18
 */
public class ZipkinMysqlStatementInterceptor implements QueryInterceptor {

    protected final Logger logger = LoggerFactory.getLogger(getClass());
    /**
     * 拦截sql的时间,超过该时间的SQL即进行预警等操作 单位 毫秒
     */
    public static final Integer MAX_SQL_EXECUTE_TIME = 3000;

    /**
     * 执行sql的开始时间
     */
    private final ThreadLocal<Long> startTimestamp = new ThreadLocal<>();

    public <T extends Resultset> T preProcess(Supplier<String> sqlSupplier, Query interceptedQuery) {
        startTimestamp.set(System.currentTimeMillis());
        return superPreProcess(sqlSupplier, interceptedQuery);
    }

    public <T extends Resultset> T postProcess(Supplier<String> sql, Query interceptedQuery, T originalResultSet, ServerSession serverSession) {

        Long startTime = startTimestamp.get();
        startTimestamp.remove();
        if (startTime != null) {
            //SQL执行总时长
            long time = System.currentTimeMillis() - startTime;
            if (time >= MAX_SQL_EXECUTE_TIME) {

//                String traceId = LoggerTraceMcUtil.getSleuthTraceId();
//                String msg = String.format("总耗时:{%d}毫秒,traceId={%s},执行的SQL\n\n%s", time, traceId,sql.get());
//                DingDingUtil.sendDDMessage("https://oapi.dingtalk.com/robot/send?access_token=7fb3d4640752b0fd1f44bea1b40a9553539ca62fc62af2e991b85ee10ac17d5f",msg,false);
            }
        }
        T t = superPostProcess(sql, interceptedQuery, originalResultSet, serverSession);
        return t;
    }
    private MysqlConnection connection;
    private boolean interceptingExceptions;

    public <T extends Resultset> T superPreProcess(Supplier<String> sqlSupplier, Query interceptedQuery) {
        Span span = ThreadLocalSpan.CURRENT_TRACER.next();
        if (span != null && !span.isNoop()) {
            String sql = (String)sqlSupplier.get();
            int spaceIndex = sql.indexOf(32);
            span.kind(Span.Kind.CLIENT).name(spaceIndex == -1 ? sql : sql.substring(0, spaceIndex));
            span.tag("sql.query", sql);
            parseServerIpAndPort(this.connection, span);
            span.start();
            return null;
        } else {
            return null;
        }
    }

    public <T extends Resultset> T superPostProcess(Supplier<String> sql, Query interceptedQuery, T originalResultSet, ServerSession serverSession) {
        if (this.interceptingExceptions && originalResultSet == null) {
            return null;
        } else {
            Span span = ThreadLocalSpan.CURRENT_TRACER.remove();
            if (span != null && !span.isNoop()) {
                span.finish();
                return null;
            } else {
                return null;
            }
        }
    }

    static void parseServerIpAndPort(MysqlConnection connection, Span span) {
        try {
            URI url = URI.create(connection.getURL().substring(5));
            String remoteServiceName = connection.getProperties().getProperty("zipkinServiceName");
            String host;
            if (remoteServiceName == null || "".equals(remoteServiceName)) {
                host = getDatabaseName(connection);
                if (host != null && !host.isEmpty()) {
                    remoteServiceName = "mysql-" + host;
                } else {
                    remoteServiceName = "mysql";
                }
            }

            span.remoteServiceName(remoteServiceName);
            host = getHost(connection);
            if (host != null) {
                span.remoteIpAndPort(host, url.getPort() == -1 ? 3306 : url.getPort());
            }
        } catch (Exception var5) {
        }

    }

    private static String getDatabaseName(MysqlConnection connection) throws SQLException {
        return connection instanceof JdbcConnection ? ((JdbcConnection)connection).getCatalog() : "";
    }

    private static String getHost(MysqlConnection connection) {
        return !(connection instanceof JdbcConnection) ? null : ((JdbcConnection)connection).getHost();
    }

    public boolean executeTopLevelOnly() {
        return true;
    }

    public QueryInterceptor init(MysqlConnection mysqlConnection, Properties properties, Log log) {
        String exceptionInterceptors = properties.getProperty("exceptionInterceptors");
        ZipkinMysqlStatementInterceptor interceptor = new ZipkinMysqlStatementInterceptor();
        interceptor.connection = mysqlConnection;
        interceptor.interceptingExceptions = exceptionInterceptors != null && exceptionInterceptors.contains(ZipkinTracingExceptionInterceptor.class.getName());
        if (!interceptor.interceptingExceptions) {
            log.logWarn("TracingExceptionInterceptor not enabled. It is highly recommended to enable it for error logging to Zipkin.");
        }

        return interceptor;
    }

    public void destroy() {
    }
}